import argparse
import yaml
import time
from hexdump import hexdump
from serial import Serial
from tqdm import tqdm


def load_config(yaml_file):
    try:
        with open(yaml_file, 'r') as file:
            data = yaml.safe_load(file)
            return data
    except:
        return None

def data_generate(cfg):
    datas = []
    for d in range(cfg['size']):
        datas.append(d & 0xff)
    return bytes(datas)

def loop_data_check(cfg, data):
    if len(data) != cfg['size']:
        return False
    for i in range(cfg['size']):
        if data[i] != (i & 0xff):
            return False
    return True

def rx_data_check(cfg, data):
    for i in range(cfg['size']):
        if data[i] != ((data[0] + i) & 0xff):
            return False
    return True

def sp_loop_test(cfg, serial):
    txdata = data_generate(cfg)
    if cfg['hexdump']:
        print('tx dump:')
        hexdump(txdata)
    serial.write(txdata)
    rxdata = serial.read(cfg['size'])
    if not rxdata:
        return False
    if cfg['hexdump']:
        print('rx dump:')
        hexdump(rxdata)
    return loop_data_check(cfg, rxdata)
    
def sp_tx_test(cfg, serial):
    txdata = data_generate(cfg)
    if cfg['hexdump']:
        print('tx dump:')
        hexdump(txdata)
    serial.write(txdata)
    time.sleep(cfg['tx_interval'])
    return True

def sp_rx_test(cfg, serial):
    rxdata = serial.read(cfg['size'])
    if cfg['hexdump']:
        print('rx dump:')
        hexdump(rxdata)
    ret = rx_data_check(cfg, rxdata)
    if ret != True:
        print('rx dump:')
        hexdump(rxdata)
        return ret
    return True

def sp_port_test(cfg, com):
    try:
        serial = Serial(port=com, baudrate=cfg['baudrate'], timeout=1)
        for i in tqdm(range(cfg['epochs'])):
            if cfg['mode'] == 'loop':
                ret = sp_loop_test(cfg, serial)
            elif cfg['mode'] == 'tx':
                ret = sp_tx_test(cfg, serial)
            elif cfg['mode'] == 'rx':
                ret = sp_rx_test(cfg, serial)
            if ret != True:
                return False
        return True
    except Exception as e:
        print(e)

def dp_loop_test(cfg, serial):
    txdata = data_generate(cfg)
    if cfg['hexdump']:
        print('tx dump:')
        hexdump(txdata)
    serial[0].write(txdata)
    rxdata = serial[1].read(cfg['size'])
    if not rxdata:
        return False
    if cfg['hexdump']:
        print('rx dump:')
        hexdump(rxdata)
    return loop_data_check(cfg, rxdata)
    
def dp_pair_test(cfg, com):
    try:
        serial_tx = Serial(port=com[0], baudrate=cfg['baudrate'], timeout=1)
        serial_rx = Serial(port=com[1], baudrate=cfg['baudrate'], timeout=1)
        for i in range(cfg['epochs']):
            ret = dp_loop_test(cfg, (serial_tx, serial_rx))
            if ret != True:
                return ret
        return True
    except Exception as e:
        print(e)
        return False

# Double port test
def dp_test(cfg):
    num_ports = len(cfg['ports'])
    if num_ports % 2 != 0:
        print('ports must multiple of 2')
        return
    if cfg['mode'] != 'loop':
        print('double port must be loop mode')
        return
    num_pairs = int(num_ports / 2)
    for i in range(num_pairs):
        com_tx = cfg['ports'][i*2]
        com_rx = cfg['ports'][i*2 + 1]
        print(com_tx)
        print(com_rx)
        ret = dp_pair_test(cfg, (com_tx, com_rx))
        if ret != True:
            print(f'Testing ({com_tx},{com_rx}) Error')
        else:
            print(f'Testing ({com_tx},{com_rx}) Success')
    
# Signle port test
def sp_test(cfg):
    for port in cfg['ports']:
        ret = sp_port_test(cfg, port)
        if ret != True:
            print(f'Testing {port} Error')
        else:
            print(f'Testing {port} Success')

def main():
    cfg = load_config('config.yaml')
    if cfg['double_port'] == True:
        dp_test(cfg)
    elif cfg['double_port'] == False:
        sp_test(cfg)

if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        print(e)